import sys
import os
import warnings
import logging
from contextlib import contextmanager
import tqdm
from functools import wraps
from pathlib import Path
from queue import Queue
from threading import Event
import queue
import time

from get_filenames import GetFilename


class CosyVoiceWrapper:
    def __init__(
            self, 
            # event什么的
            is_recording_event= Event(),
            tts_stopped_event= Event(),
            exit_event= Event(),
            # 队列
            tts_text_queue: Queue = Queue(),
            playmusic_queue: Queue = Queue(),

            model_path='CosyVoice2/pretrained_models/CosyVoice2-0.5B', 
            show_logs=True, 
            show_warnings=True, 
            show_print=True,
            show_tqdm=True,
            matcha_tts_path='CosyVoice2/third_party/Matcha-TTS',
            default_text='',
            default_text_path='assets/texts/test.txt',
            prompt_text='',
            prompt_path='assets/prompts',
            prompt_text_file='康辉01.txt',
            prompt_audio_file='康辉01.wav',
            delete_audio=False
            ):
        """
        如果直接输入text，则忽略text_path内容
        如果直接输入prompt_text，则忽略prompt_text_file内容，不过感觉prompt_text完全用不上
        """
        # 事件
        self.is_recording_event = is_recording_event
        self.tts_stopped_event = tts_stopped_event
        self.exit_event = exit_event
        # 队列
        self.tts_text_queue = tts_text_queue
        self.playmusic_queue = playmusic_queue
        # 其他状态
        self.was_interrupted = False
        self.change_prompt = False

        self.model_path = model_path
        self.show_logs = show_logs
        self.show_warnings = show_warnings
        self.show_print = show_print  # 新增参数控制print
        self.show_tqdm = show_tqdm  # 新增参数控制tqdm
        self.matcha_tts_path = matcha_tts_path
        self.default_text = default_text
        self.default_text_path = default_text_path
        self.prompt_text = prompt_text
        # 找到prompt文件夹
        self.prompt_path = prompt_path
        self.prompt_text_file = prompt_text_file
        self.prompt_audio_file = prompt_audio_file
        self.get_filename = GetFilename("temps/cosyvoice_", ".wav")
        self.delete_audio = delete_audio
        # 所有的生成的文件
        self.filenames = []
        # 获取prompts文件夹下的所有文件
        self.prompt_files = []
        # 可用的prompts
        self.prompts = []
        
        # 根据参数配置日志和警告
        self._configure_logging()
        self._configure_warnings()

        # 初始化prompts
        self._initialize_prompt_files()

        # 防止烦人的警告
        class IgnoreSpecificWarning(logging.Filter):
            def filter(self, record):
                # 检查日志消息是否包含特定内容
                return "too short than prompt text" not in record.getMessage()

        # 给根日志器添加过滤器
        logging.getLogger().addFilter(IgnoreSpecificWarning())

        # 禁用print
        if not self.show_logs:
           # 保存原始stdout和stderr
            self.original_stdout = sys.stdout
            self.original_stderr = sys.stderr  # 新增：保存原始stderr
            # 重定向到空设备
            null_device = open(os.devnull, 'w')
            sys.stdout = null_device
            sys.stderr = null_device  # 新增：重定向stderr

        # 导入并初始化模型
        self.cosyvoice = self._initialize_model()
        # 初始化提示词文件列表和提示词列表
        self._initialize_prompt()

        # 恢复print
        if not self.show_logs:
            if hasattr(self, 'original_stdout'):
                sys.stdout = self.original_stdout
            if hasattr(self, 'original_stderr'):  # 新增：恢复stderr
                sys.stderr = self.original_stderr
            # 关闭空设备文件（可选，避免资源泄漏）
            if hasattr(self, 'null_device'):
                self.null_device.close()

    def _initialize_prompt_files(self):
        """
        读取文件夹内的所有prompt文件
        """
        foder_path = Path(self.prompt_path)
        self.prompt_files = [f.name for f in foder_path.iterdir() if f.is_file()]   
        # 找到所有txt文件的文件名
        txt_files = set()
        for file in self.prompt_files:
            if file.lower().endswith('.txt'):
                core_name = file[:-4].strip()
                txt_files.add(core_name)
        
        # 检查是否有对应的wav文件,并加入prompts列表
        for file in self.prompt_files:
            if file.lower().endswith('.wav'):
                core_name = file[:-4].strip()
                if core_name in txt_files:
                    self.prompts.append(core_name)
                    
    def _configure_logging(self):
        """豆包写的"""
        if not self.show_logs:
            # 禁用所有日志
            logging.basicConfig(level=logging.CRITICAL + 1)
    
    def _configure_warnings(self):
        """配置警告输出"""
        if not self.show_warnings:
            # 忽略所有警告
            warnings.filterwarnings("ignore")
        else:
            # 显示所有警告
            warnings.filterwarnings("default")

    def _initialize_prompt(self):
        """
        读取一个prompt的txt文件
        """
        if self.prompt_text == '':
            with open(self.prompt_path + "/" + self.prompt_text_file, 'r', encoding='utf-8') as f:
                self.prompt_text = f.read()
    
    def _initialize_model(self):
        """初始化CosyVoice模型"""
        sys.path.append(self.matcha_tts_path)
        from cosyvoice.cli.cosyvoice import CosyVoice2
        
        return CosyVoice2(self.model_path,load_jit=False,load_trt=False,load_vllm=False,fp16=False)
    
    @contextmanager
    def _tqdm_controls(self):
        """
        添加一个上下文管理器用于临时控制tqdm
        完全看不懂这这段代码，豆包写的
        """
        if not self.show_tqdm:
            # 保存原始的tqdm构造函数
            original_init = tqdm.tqdm.__init__
            
            # 定义新的构造函数，强制禁用进度条
            @wraps(original_init)
            def new_init(self, *args, **kwargs):
                # 无论传入什么参数，强制设置disable=True
                kwargs['disable'] = True
                original_init(self, *args, **kwargs)
            
            # 替换tqdm的构造函数
            tqdm.tqdm.__init__ = new_init
            
            try:
                yield
            finally:
                # 恢复原始构造函数
                tqdm.tqdm.__init__ = original_init
        else:
            yield

    def generate_audio(
            self, 
            text='',
            text_path=''
            ):
        """生成语音并保存为WAV文件"""
        from cosyvoice.utils.file_utils import load_wav
        import torchaudio
        
        # 加载提示音频
        prompt_speech_16k = load_wav(self.prompt_path + "/" + self.prompt_audio_file, 16000)

        if text_path == '':
            text_path = self.default_text_path
        
        if text == '':
            with open(text_path, 'r', encoding='utf-8') as f:
                text = f.read()

        with self._tqdm_controls():
            # 生成语音
            results = self.cosyvoice.inference_zero_shot(
                text,
                self.prompt_text,
                prompt_speech_16k
            )
            # 保存生成的音频
            output_files = []
            for i, result in enumerate(results):
                output_path = self.get_filename.get_filename()
                torchaudio.save(output_path, result['tts_speech'], self.cosyvoice.sample_rate)
                output_files.append(output_path)
                
                # 只有允许print时才输出
                if self.show_print:
                    print(f"已保存音频文件: {output_path}")
                self.filenames.append(output_path)
        return output_files
    
    def generate_audio_loop(self):
        """多次运用版，主循环"""
        from cosyvoice.utils.file_utils import load_wav
        import torchaudio

        # 冒充初始化完成标志
        self.tts_stopped_event.set()
        print("✅ 文字转语音模型 \t| CosyVoice2 \t→ 加载完成")
        
        while not self.exit_event.is_set():
            if self.is_recording_event.is_set():
                if not self.was_interrupted:
                    self._clear_queue(self.tts_text_queue)
                    self.tts_stopped_event.set()
                    self.was_interrupted = True
                time.sleep(0.02)
                continue
            else:
                self.was_interrupted = False

            try:
                # 循环等待时间
                text = self.tts_text_queue.get(timeout=0.02)

                # 再次检查录音状态
                if self.is_recording_event.is_set():
                    self._clear_queue(self.tts_text_queue)
                    self.tts_stopped_event.set()
                    self.was_interrupted = True
                    continue

                # 文字转语音，加载提示音频
                prompt_speech_16k = load_wav(self.prompt_path + "/" + self.prompt_audio_file, 16000)

                with self._tqdm_controls():
                    # 生成语音
                    results = self.cosyvoice.inference_zero_shot(
                        text,
                        self.prompt_text,
                        prompt_speech_16k
                    )
                    # 保存生成的音频
                    output_files = []
                    for i, result in enumerate(results):
                        output_path = self.get_filename.get_filename()
                        torchaudio.save(output_path, result['tts_speech'], self.cosyvoice.sample_rate)
                        output_files.append(output_path)
                        # 只有允许print时才输出
                        if self.show_print:
                            print(f"已保存音频文件: {output_path}")
                        self.filenames.append(output_path)
                        # 加入队列,再次加入检测是否正在录音
                        if not self.is_recording_event.is_set():
                            self.playmusic_queue.put(output_path)
                    self.tts_text_queue.task_done()

            except queue.Empty:
                continue
        # 防止文件被占用
        self._delete_files()
    
    def _delete_files(self):
        if self.delete_audio:
            for file in self.filenames:
                for _ in range(3):  # 最多重试次
                    try:
                        os.remove(file)
                        break  # 删除成功就退出重试循环
                    except PermissionError:
                        time.sleep(1)  # 等1秒再试
                else:
                    # 5次都失败
                    print(f"文件 {file} 无法删除，可能被占用")

    def _clear_queue(self, queue: Queue):
        """清空队列中的所有音频文件"""
        while not queue.empty():
            try:
                queue.get_nowait()
                queue.task_done()
            except queue.Empty:
                break

    def start(self):
        self.generate_audio_loop()

# 使用示例
if __name__ == "__main__":
    pass

